#!/usr/bin/env python

import readline, sys, copy, getopt, collections, re, string
import cryptanal
from operator import itemgetter
import rotor, vigenere, substitution, freqanalysis, adfgvx

"""
The shell class is the front end for cryptanal
and includes interfaces to the entire backend
"""
class shell:
  def __init__(self, cryptobj=None):
    print"""
WWW         WW eEeEeEeE LL        CCCCC    OOOO    MMMM    MMMM  eEeEeEeE
 WW         W  EE       LL       Cc      OOO  OOO  MM MM  M  MM  EE
  WW       WW  EeEeE    LL      CC       OO    OO  MM  MMM   MM  EeEeE
  WWw WW  WW   EE       LL       Cc      OOO  OOO  MM        MM  EE
   WWW  WWW    eEeEeEeE LlLlLlL   CCCCC    OOOO    MM        MM  eEeEeEeE

                             TO CRYPTO-SHELL
  (Useful for deciphering what little Susie is writing to little Billy)
"""
    self.rootcrypto = cryptobj
    self.currcrypto = copy.deepcopy(cryptobj)

    #setup the tab completion information here
    self.commands = ["help", "print", "save", "revert", "open", "sub", "trans", "quit", "enigma", "vigenere", "substitution", "load", "analyze", "suggestsub", "adfgvx"]
    readline.set_completer(self.completer)
    readline.parse_and_bind("tab: complete")

  ############################################
  # Below is the main front-end interface
  # including support for gnu readline
  # and print functions
  ############################################

  ############################################
  # The function for tab-completion
  ############################################
  def completer(self, word, index):
    matches = [c for c in self.commands if c.startswith(word)]
    try:
      return matches[index] + " "
    except IndexError:
      pass

  ############################################
  # The mainloop
  ############################################
  def mainloop(self):
    while 1:
      try:
        command=raw_input('> ').lstrip()
        cmdl = command.lower() # save typing
        if not len(cmdl.strip()): # ignore press enter on blank line
          pass
        elif command.lower().startswith('help'):
          self.help(command[4:].lstrip())
        elif command.lower().startswith('save'):
          self.save(command[4:].lstrip())
        elif command.lower().startswith('revert'):
          self.Revert(command[6:].strip())
        elif command.lower().startswith('print'):
          self.Print(command[5:].strip())
        elif command.lower().startswith('substitution'):
          args = command[12:].split()
          self.substitution(args)
        elif command.lower().startswith('sub'):
          args = command[3:].split()
          self.sub(args)
        elif command.lower().startswith('enigma'):
          args = command[6:].split()
          self.enigma(args)
        elif command.lower().startswith('vigenere'):
          args = command[8:].split()
          self.vigenere(args)
        elif command.lower().startswith('adfgvx'):
          args = command[6:].split()
          self.adfgvx(args)
        elif cmdl.startswith('analyze'):
          args = command[7:].split()
          self.analyze(args)
        elif cmdl.startswith('suggestsub'):
          args = command[10:].split()
          if args:
            self.suggestsub(args[0])
          else:
            print "No language specified. Using English."
            self.suggestsub("english")
        elif command.lower().startswith('trans'):
          args = command[5:].split()
          self.transpose(args)
        elif command.lower().startswith('load') or command.lower().startswith('open'):
          args = command[4:].split()
          self.load(args)
        elif command.strip().lower() == "quit" or command.strip().lower() == "exit":
          print "come again soon...  "
          sys.exit()
        else:
          print "Error: command not recognized: type 'help' for usage"
      except KeyboardInterrupt:
        print "come again soon... "
        sys.exit()

  ############################################
  # This is the main help function
  ############################################
  def help(self, args):
    print "The following commands are available:"
    for x in self.commands:
      print " ", x
    print "Try '<command> --help' for usage on a specific command."
    print

  ############################################
  # This is a function for reverting back to
  # the original text
  ############################################
  def Revert(self, args):
    self.currcrypto = copy.deepcopy(self.rootcrypto)

  ############################################
  # This is the print function which prints
  # the cipher in each of the various ways
  ############################################
  def Print(self, args):
    printhelp = 0
    try:
      if args.strip() == "current":
        self.currcrypto.printcrypt()
      elif args.strip() == "orig":
        self.rootcrypto.printcrypt()
      elif args.strip() == "freq":
        self.currcrypto.printsorteddict()
      elif args.strip() == "graph":
        self.PrintGraph()
      elif args.strip() == "help":
        self.PrintUsage()
      else:
        self.PrintUsage()
    except AttributeError:
      print "Please load a file first by typing \"load <file-name>\""

  ############################################
  # This is a function for printing graphs.
  # Right now, it is overly-simplified.
  # The only way to print is the cipher vs
  # english, grouped by letter. More options
  # will be added later.
  ############################################
  def PrintGraph(self):
    try:
      import freqanalysis, os, signal, pickle
      newpid = os.fork()
      #zombies must be shot in the head
      def zombiekiller(a, b):
        os.wait()
      signal.signal(signal.SIGCHLD, zombiekiller)
      #fork a new process so the graph doesn't tie up the shell
      if newpid == 0:
        #m = freqanalysis.xfrequencyTable(self.currcrypto.getfreq())
        #m.printEnglishvFreq()
        try:
          os.execv("./freqanalysis.py", ["freqanalysis.py", pickle.dumps(self.currcrypto.getfreq())])
          raise SystemExit
        except AttributeError:
          print "Please load a file first by typing \"load <file-name>\""
    except SystemExit:
       sys.exit()
    except:
      print "Error: check that numpy and matplotlib are installed"
      print "you also probably can't use windows because there is a fork"
      print sys.exc_info()[0]

  ############################################
  # Usage for the different print options
  ############################################
  def PrintUsage(self):
    print """Usage: print [current | orig | freq]

  current		prints current crypto decoding
  orig			prints original ciphertext
  freq			prints a simple frequency analyses count
  graph 		prints various graphs - must have matplotlib installed

  help			print this screen
  """

  ############################################
  # Below are the Substitution an Transpose
  # interfaces which are common to many
  # cryptographic systems
  ###########################################
  def sub(self, args):
    #set default argumetns
    subhelp = False
    startrange = 0
    endrange = None
    jumpval = 1
    try:
      fromstr = args[0]
      tostr = args[1]
    #get arguments
      options, arguments = getopt.getopt(args[2:], "s:e:j:", ["startrange=", "endrange=", "jumpval="])
      for o, a in options:
        if o in ("-s", "--startrange"):
          startrange = int(a)
        elif o in ("-e", "--endrange"):
          endrange = int(a)
        elif o in ("-j", "--jumpval"):
          jumpval = int(a)
        else:
          subhelp = True
          break
      if subhelp == True:
        self.subUsage()
      else:
        try:
          self.currcrypto.substitute(fromstr, tostr, startrange, endrange, jumpval)
        except AttributeError:
          print "Please load a file first by typing \"load <file-name>\""
    except:
      self.subUsage()
      return

  ############################################
  # Usage for the substitution commands
  ############################################
  def subUsage(self):
    print """Usage: sub fromstring tostring [-s | --startrange startrange] [-e | --endrange endrange] [-j | --jumpval jumpvalue]

  fromstring		original character to be replaced
  tostring		character to replace fromstring
  -s,--startrange	index (starting with 0) in which to start replacing letters
  -e,--endrange		index in which to stop replacing letters
  -j,--jumpvalue	allows operation on series of values that are not necessarily next to eachother"""

  def transpose(self, args):
    #set default arguments
    transHelp = False
    shiftl = 1
    shiftr = 0
    srange = 0
    shifterrorcheck = 0
    try:
      endrange = len(self.currcrypto.letterlist) - 1
      options, arguments = getopt.getopt(args, "l:r:s:e:j:", ["lshift=", "rshift=", "srange=", "endrange=", "help"])
      for o, a in options:
        if o in ("-l", "--lshift"):
          shiftl = int(a)
          #this is only an error if right shift is also set
          shifterrorcheck += 1
        elif o in ("-r", "--rshift"):
          shiftr = int(a)
          #only an error if shiftl is also set
          shifterrorcheck += 1
        elif o in ("-s", "--srange"):
          srange = int(a)
        elif o in ("-e", "--endrange"):
          endrange = int(a)
        elif o in ("--help"):
          transHelp = True
        else:
          transHelp = True
        if transHelp == True:
          break
      if shifterrorcheck >= 2 or arguments:
        transHelp = True
      if transHelp == True:
        self.transHelp()
      else:
        if shiftr != 0:
          shiftl = len(self.currcrypto.letterlist) - shiftr
        self.currcrypto.transpose(shiftl, srange, endrange)
    except AttributeError:
      print "Please load a file first by typing \"load <file-name>\""
    except:
      self.transHelp()
      return

  ############################################
  # Usage for the transpose commands
  ############################################
  def transHelp(self):
    print """Usage: trans [[-l --lshift value | -r --rshift value] [-s | --srange] [-e | --endrange] | --help]

  -l, --lshift		number of places to shift left
  -r, --rshift		number of places to shift right
  -s, --srange		where to start transposing
  -e, --endrange	where to end transposing
  --help		print this message and exit
  """

  ###########################################
  # This is a function for saving what
  # has been done
  ###########################################
  def save(self, name):
    if name is "":
      name = "out.txt"
    f=open(name, 'w')
    try:
      f.write(self.currcrypto.getcrypt())
    except AttributeError:
      print "Please load a file first by typing \"load <file-name>\""
    f.close()
    print "Saved in " + name

  ###########################################
  # This is a function for loading a new file
  ###########################################
  def load(self, args):
    try:
      letterlist = open(args[0]).read()
      rootphrase = cryptanal.CryptAnal(letterlist)
      self.rootcrypto = rootphrase
      self.currcrypto = copy.deepcopy(rootphrase)
      print args[0] + " opened\n"
    except IndexError:
      print "> Error loading file"

  ###########################################
  # Front end for Substitution Cipher
  ###########################################
  def substitution(self, args):
    try:
      cryptToSend = self.currcrypto.getcrypt()
      cryptToSend = cryptToSend.replace(' ', '')  # deleting all white space
#    cryptToSend = cryptToSend.lower() # make all letters lower-case for comparison with english letter frequencies
      sub = substitution.Substitution(cryptToSend)
      cipherFreqTable = sub.getFreqTable()
      freq = freqanalysis.xfrequencyTable(cipherFreqTable)
      myNewText = freq.replaceAll()
      sub.mainLoop()
#       freqComparison = sub.getFrequency()
    except AttributeError:
      print "Please load a file first by typing \"load <file-name>\""

  ###########################################
  # Front end for Vigenere Cipher
  ###########################################
  def vigenere(self, args):
    usageProblem = False
    try:
      options, arguments = getopt.getopt(args, "k:edh", ["key=", "encrypt", "decrypt", "help"])
    except getopt.GetoptError, err:
      usageProblem = True
    if usageProblem == False:
      key = "default"
      encrypt = 0
      decrypt = 0
      doVigenere = False
      help = 0
      for o, a in options:
	if o in ("-e", "--encrypt"):
	  doVigenere = True
	  encrypt = 1
	elif o in ("-d", "--decrypt"):
	  doVigenere = True
	  decrypt = 1
	elif o in ("-k", "--key"):
          key = a
        elif o in ("-h", "--help"):
          help = 1
      if help == 1:
        self.vigUsage()
      if doVigenere is True:
        try:
	  vig = vigenere.Vigenere(self.currcrypto.getcrypt().strip(), key)
	  if encrypt == 1 and decrypt == 1:
	    #usage problem, -d and -e cannot be used at the same time
            self.vigUsage()
	  elif encrypt == 1:
	    res = vig.cypher()
	    self.currcrypto.setcrypt(res)
	  elif decrypt == 1:
	    res = vig.decypher()
	    self.currcrypto.setcrypt(res)
        except AttributeError:
          print "Please load a file first by typing \"load <file-name>\""
      else:
        #at least -d or -e have to be used
        self.vigUsage()

  ###########################################
  # Usage for vigenere commands
  ###########################################
  def vigUsage(self):
    print """Usage: vigenere [-e | --encrypt] | [-d | --decrypt] [-h | --help] [-k | --key cipherKey] "
    -e, --encrypt                Encrypt input.  Can not be used with the -d switch.
    -d, --decrypt                Decrypt input.  Can not be used with the -e switch.
    -k, --key                    Set the key used to encrypt or decrypt an input (string). If not used, key="default"
    -h, --help                   Print this screen.
    """

  def adfgvx(self, args):
    usageProblem = False
    try:
      options, arguments = getopt.getopt(args, "t:edh", ["transposekey=", "encrypt", "decrypt", "help"])
    except getopt.GetoptError, err:
      usageProblem = True
    if usageProblem == False:
      transposekey = "default"
      encrypt = 0
      decrypt = 0
      doADFGVX = False
      help = 0
      for o, a in options:
        if o in ("-e", "--encrypt"):
          doADFGVX = True
          encrypt = 1
        elif o in ("-d", "--decrypt"):
          doADFGVX = True
          decrypt = 1
        elif o in ("-t", "--transposekey"):
          transposekey = a
        elif o in ("-h", "--help"):
          help = 1
      if help == 1:
        self.ADFGVXUsage()
      if doADFGVX is True:
        try:
          myADFGVX = adfgvx.ADFGVX(self.currcrypto.getcrypt().strip(), transposekey)
          if encrypt == 1 and decrypt == 1:
            self.ADFGVXUsage()
          elif encrypt == 1:
            res = myADFGVX.encrypt()
            self.currcrypto.setcrypt(res)
          elif decrypt == 1:
            res = myADFGVX.decrypt()
            self.currcrypto.setcrypt(res)
        except AttributeError:
          print "Please load a file first by typing \"load <file-name>\""
      else:
        self.ADFGVXUsage()

  ###########################################
  # Usage for ADFGVX commands
  ###########################################
  def ADFGVXUsage(self):
    print """Usage: adfgvx [-e | --encrypt] | [-d | --decrypt] [-h | --help] [-t | --transposekey Key] "
    -e, --encrypt                Encrypt input.  Can not be used with the -d switch.
    -d, --decrypt                Decrypt input.  Can not be used with the -e switch.
    -t, --transposekey           Set the key used to encrypt or decrypt an input (string). If not used, transposekey="default"
    -h, --help                   Print this screen.
    """

  ###########################################
  # Front end for Enigma cipher
  ###########################################
  def enigma(self, args):
    # Default values needed for syntax checking and -h options.
    encipher = 0
    decipher = 0
    enigmaU = False
    # The next section parses the arguments "args" grabed from the command line and passed down from the shell.
    try:
      options, arguments = getopt.getopt(args, "r:k:eduh", ["key=", "rotors=", "encrypt", "decrypt", "unknown", "help"])
    except getopt.GetoptError, err:
      # this exception is executed if there is an option passed that is not defined.  Normally this would have an exit,
      # or there would be an exit in the usage and you would call the usage here.
      enigmaU = True
    # Since Python does not have a goto command and I don't want to exit the shell I set a variable that makes most of
    # the code get skipped if there is a usage problem.
    if enigmaU is not True:
      # Default settings for the cipher program.
      keys = "default"
      kVar = False
      rotors = 3
      rVar = False
      encryption = True
      unknown = False
      # conditions and variables that are set depending on the options passed
      for o, a in options:
        if o in ("-d", "--decrypt"):
          encryption = False
          decipher = 1
        elif o in ("-e", "--encrypt"):
          encryption = True
          encipher = 1
        elif o in ("-h", "--help"):
          enigmaU = True
        elif o in ("-k", "--key"):
          keys = a
          kVar = True
        elif o in ("-r", "--rotors"):
          rotors = a
          rVar = True
        elif o in ("-u", "--unknown"):
          encryption = False
          decipher = 1
          unknown = True
          if rVar is True and kVar is True:
            enigmaU = True
        elif enigmaU is True:
          # this conditional breaks out of the for loop if the help options is called.
          break
        else:
          # this conditional sets the usage flag and breaks out of the for loop if there is an unhandled condtion
          # in the command line.
          enigmaU = True
          break
      if encipher and decipher:
        # this conditional flags if the -e and -d are both called.  They are mutually exclusive.
        print "The -e and -d/u are mutually exclusive.  "
        enigmaU = True
      try:
        rotors = int(rotors)
      except:
        enigmaU = True
        print "Rotors must be an int.  "
      if unknown is True:
        # print "unknown switch set.  "
        if rVar is True:
          enigmaU = True
          print "The key (-k) must be specified, and the rotor (-r) must not be specified.  "

    if enigmaU is True:
      self.enigmaUsage()
    else:
      rt = rotor.newrotor(keys, rotors)
      if encryption == True:
        try:
          self.currcrypto.setcrypt(rt.encrypt(self.currcrypto.getcrypt()))
          print "Key and rotors are set to: " + keys + ", " + str(rotors)
        except AttributeError:
          print "Please load a file first by typing \"load <file-name>\""
      else:
        if unknown == False:
          try:
            self.currcrypto.setcrypt(rt.decrypt(self.currcrypto.getcrypt()))
            print "Key and rotors are set to: " + keys + ", " + str(rotors)
          except AttributeError:
            print "Please load a file first by typing \"load <file-name>\""
        else:
          # This loop will need to be run when searching for a decryptions configuration
          # loop through more than one word in a string.
          percent = 0.0
          decriptIterator = 0
          while percent < 50 and decriptIterator < 20:
            percent = 0.0
            percentNum = 0
            percentDen = 0
            rt = rotor.newrotor(keys, decriptIterator)
            # self.currcrypto.setcrypt(rt.decrypt(self.currcrypto.getcrypt()))
            try:
              for word in rt.decrypt(self.currcrypto.getcrypt()).split():
                if dictionary([word.lower().rstrip('.')]):
                  percentNum += 1
                percentDen += 1
              percent = (float(percentNum)/float(percentDen))*100
              if percent > 50:
                self.currcrypto.setcrypt(rt.decrypt(self.currcrypto.getcrypt()))
                print "Key and rotors are set to: " + keys + ", " + str(decriptIterator)
              decriptIterator += 1
            except AttributeError:
              print "Please load a file first by typing \"load <file-name>\""

  ###########################################
  # Usage for enigma
  ###########################################
  def enigmaUsage(self):
    print """Usage: enigma [-e | --encrypt] | [-d | --decrypt] [-u | --unknown] [-h | --help] [-k | --key cipherKey] [-r | --rotor cipherRotorNumber] "
    -e, --encrypt    Encrypt input.  Can not be used with the -d switch.
    -d, --decrypt    Decrypt input.  Can not be used with the -e switch.
    -u, --unknown    Unknown rotor number for the given encrypted input.
                     Note this option only is used for decryption it has
                     no effect if the -e or --encrypt switchs are set.
    -k, --key        Set the key used to encrypt or decrypt an input (string).
    -r, --rotors     Set the number encryption/decryption rotors (int).
    -h, --help       Print this screen.
  """

  ###########################################
  # Entrypoint for command: analyze
  # This should eventually be split up into several classes
  # and do a lot more. Right now, it only analyzes single-character substitution.
  # It also makes several assumptions, which may not be correct:
  # 1) Only letters are substituted
  # 2) All letters are substituted by only letters (but they can be any case)
  #     Spaces, punctuation, numbers, etc. are ignored
  # 3) The crypted data is large enough to do a meaningful frequency analysis.
  # 4) The crypted file is in a language which uses a Roman alphabet.
  #
  # Obviously there's a nearly unlimited amount of functionality that could be
  # added to the encryption analysis.
  ###########################################
  def analyze(self, args):
    if args:
      self.analyzeUsage()

    freq_current = list()
    dict = self.currcrypto.getfreq()
    for key in dict.keys():
      if key.lower() in string.lowercase:
        freq_current.append(dict[key])
    freq_current.sort(reverse=True)

    # The key used to be the language code (e.g. 'en', 'tr', etc.)
    # but it turns out the key is convenient for on-screen display.
    freq_dict = {
      'English' : self.getDictAsList(freqanalysis.EnglishFreqTable),
      'French' : self.getDictAsList(freqanalysis.FrenchFreqTable),
      'German' : self.getDictAsList(freqanalysis.GermanFreqTable),
      'Spanish' : self.getDictAsList(freqanalysis.SpanishFreqTable),
      'Italian' : self.getDictAsList(freqanalysis.ItalianFreqTable),
      'Turkish' : self.getDictAsList(freqanalysis.TurkishFreqTable),
      'Swedish' : self.getDictAsList(freqanalysis.SwedishFreqTable)
    }

    # Gets the delta (difference) between the frequency of the x'th
    # letter in the current crypted file and the x'th letter in each
    # supported language, then adds all of the deltas.
    # Theoretically, languages with a smaller delta are more likely
    # to be the language of the unencrypted file.
    delta_dict = {}
    print "Currently comparing current crypted data against: ",
    for key in freq_dict.keys():
      print key + ',',
      for a, b in zip(freq_dict[key], freq_current):
        delta_dict[key] = abs(a - b)
    print 'done.'

    sortedDelta = sorted(delta_dict.items(), key = itemgetter(1))

    # Maximum total delta for crypted text which is likely to be
    # substitution-cipher'ed text using the Roman alphabet.
    # Some example values deltas for context:
    # Crypted text is English? Delta between text and English distribution
    # is likely < 0.0002.
    # Crypted text is random, or encrypted with something strong like AES?
    # Delta is likely to be > 0.003
    # The default value 0.001 is an educated guess.
    maxProbableDelta = 0.001

    if sortedDelta[-1][1] >= maxProbableDelta:
      print """
*According to statistical analysis, current crypted data is probably not
encrypted with a substitution cipher. That being said..."""
      print """Substitution-encrypted text may be in one of the following languages
in decreasing order of likelihood:
---------------------------------------------------------------
LANGUAGE   SCORE"""

    for lang, delta in sortedDelta:
      # probabilityStrength is an inverse of delta.
      # It's just to give users a "bigger is better" number to help
      # quantify how close different languages are.
      probabilityStrength = 0.01 / delta_dict[lang]
      print lang.ljust(10), '%5.1f' % probabilityStrength

    print "Scores > 40 are likely candidates. Scores < 15 are very unlikely."
    print "---------------------------------------------------------------"


  # Suggest a single-character substitution based on relative frequency of letters
  # Arguments:
  #  langDict: Dictionary of letter:frequency in chosen language (see freqanalysis.py)
  #  limit: Display at most this many suggestions
  def suggestsub(self, lang):
     # There should be a global distionary of language frequencies somewhere
    if lang.lower().startswith('english')  : langDict = freqanalysis.EnglishFreqTable
    elif lang.lower().startswith('french')  : langDict = freqanalysis.FrenchFreqTable
    elif lang.lower().startswith('german')  : langDict = freqanalysis.GermanFreqTable
    elif lang.lower().startswith('spanish') : langDict = freqanalysis.SpanishFreqTable
    elif lang.lower().startswith('italian') : langDict = freqanalysis.ItalianFreqTable
    elif lang.lower().startswith('turkish') : langDict = freqanalysis.TurkishFreqTable
    elif lang.lower().startswith('swedish') : langDict = freqanalysis.SwedishFreqTable
    else:
      print "Invalid language selected."
      return

    cryptDict = {}
    fullDict = self.currcrypto.getfreq()
    for key in fullDict.keys():
      if key.lower() in string.lowercase:
        cryptDict[key] = fullDict[key]

    retval = {}
    sortedCrypt = sorted(cryptDict.items(), key = itemgetter(1), reverse = True)
    sortedLang = sorted(langDict.items(), key = itemgetter(1), reverse = True)
    print "Suggested substitutions (be sure to run analyze command first!):"
    for c, l in zip(sortedCrypt, sortedLang):
      retval[c] = l[0]
      print "Try replacing '" + c[0] + "' with '" + l[0] + "'"
    return retval

# Returns items (not keys) in dictionary as descending list.
# This function can be replaced with short inline Python code later.
  def getDictAsList(self, Dict):
    retval = list()
    for key in Dict.keys():
      retval.append(Dict[key])
    retval.sort(reverse = True)
    return retval






  ###########################################
  # Usage for command: analyze
  ###########################################
  def analyzeUsage(self):
    print """Usage: "analyze"
    Analyzes currently open file for hints as to its encryption.
    """

def dictionary(words):
  model = collections.defaultdict(lambda: 1)
  ALLWORDS = re.findall('[a-z]+', file('/etc/dictionaries-common/words').read().lower())
  for f in ALLWORDS:
    model[f] += 1
  return set(w for w in words if w in model)

if __name__ == '__main__':

  try:
    letterlist = open(sys.argv[1]).read()
    rootphrase = cryptanal.CryptAnal(letterlist)
    thisshell = shell(rootphrase)
  except IndexError:
    thisshell = shell()

  thisshell.mainloop()
